import torch.nn as nn
import torch.nn.functional as F


class NetClassify(nn.Module):
    def __init__(self):
        super(NetClassify, self).__init__()
        self.first = True
        # 通道数3，输入数据为3通道；6个卷积核，学习6种特征；卷积核大小5*5；最终输出层数是6
        self.conv1 = nn.Conv2d(3, 6, 5)
        # 池化层2*2
        self.pool = nn.MaxPool2d(2, 2)
        # 6：conv1的输出层数；16是conv2的输出，16种特征；卷积核大小5*5
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 16：conv2的输出层数；18*18输入矩阵的大小；800：全连接有800个神经元
        self.fc1 = nn.Linear(16 * 18 * 18, 800)
        # 800：fc1的输出层数；120：fc2的神经元个数
        self.fc2 = nn.Linear(800, 120)
        # 120：fc2的输出层数；10：类别个数
        self.fc3 = nn.Linear(120, 10)

    def print(self, tag: str, x):
        if self.first:
            print(tag, x.shape)

    def conv_flow(self, name: str, x, conv: nn.Conv2d):
        if self.first:
            print(f'======show {name}========')
        self.print('input', x)
        x = conv(x)
        self.print('conv', x)
        x = F.relu(x)
        self.print('relu', x)
        x = self.pool(x)
        self.print('pool', x)
        return x

    def forward(self, x):
        x = self.conv_flow('conv1', x, self.conv1)
        x = self.conv_flow('conv2', x, self.conv2)
        x = x.view(-1, 16 * 18 * 18)
        # relu是激活函数，把所有非零值转成零值
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        self.print('final', x)
        self.first = False
        return x
